package top.interview.string

import problem.Problem

import scala.collection.mutable.Buffer

class Q140WordBreakII extends Problem {
  def wordBreak(s: String, wordDict: List[String]): List[String] = {
    val dp = Array.ofDim[Int](s.length, s.length + 1)
    val dict = wordDict.zipWithIndex.toMap
    val compose = Buffer[(Int, Int)]()
    if (!break1(s, dict, 0, dp, compose)) Nil
    else sentences(s, compose)
  }

  def break1(s: String, wordDict: Map[String, Int], start: Int, dp: Array[Array[Int]],
             buffer: Buffer[(Int, Int)])
  : Boolean = {
    if (s.length == start) true
    else {
      ((start + 1) to s.length).map { i =>
        if (dp(start)(i) != 0) {
          dp(start)(i) == 1
        } else {
          val s0 = s.substring(start, i)
          if (wordDict.contains(s0)) {
            if (break1(s, wordDict, i, dp, buffer)) {
              buffer += ((i, start))
              dp(start)(i) = 1
            } else {
              dp(start)(i) = -1
            }
            dp(start)(i) == 1
          } else {
            dp(start)(i) = -1
            false
          }
        }
      }.reduceOption(_ || _).getOrElse(false)
    }
  }

  def sentences(sentence: String, compose: Buffer[(Int, Int)]): List[String] = {
    val m = compose.groupBy(_._1).mapValues(_.map(_._2))
    sentences(sentence, m, sentence.length).map(_.mkString(" ")).toList
  }

  def sentences(sentence: String, m: Map[Int, Buffer[Int]], end: Int): Buffer[Buffer[String]]
  = {
    if (end != 0) {
      m(end).flatMap { start =>
        val part = sentences(sentence, m, start)
        part.foreach(_ += sentence.substring(start, end))
        part
      }
    } else {
      Buffer(Buffer())
    }
  }

  override def test(): Unit = {
    val cases = Table(("string", "dict", "match"),
      ("catsanddog",
        List("cat", "cats", "and", "sand", "dog"),
        List("cats and dog", "cat sand dog")),
      ("pineapplepenapple",
        List("apple", "pen", "applepen", "pine", "pineapple"),
        List("pine apple pen apple",
          "pineapple pen apple",
          "pine applepen apple")
      ),
      ("catsandog",
        List("cats", "dog", "sand", "and", "cat"),
        Nil)
    )
    forAll(cases) { (s, dict, isMatch) =>
      wordBreak(s, dict) should contain theSameElementsAs isMatch
    }
  }
}
